import time
import torch

from .utils import Meter

try:
    from .datasets import CocoEvaluator, prepare_for_coco
except:
    pass


def train_one_epoch(model, optimizer, data_loader, device, epoch, args):
    model.train()
    A = time.time()
    for i, (image, target) in enumerate(data_loader):
        image = image.to(device)
        target = {k: v.to(device) for k, v in target.items()}
        losses = model(image, target)
        total_loss = sum(losses.values())
        total_loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    A = time.time() - A
    print("结束一个epoch，耗时{}，设备为{}".format(A, device))


from torchvision.utils import draw_segmentation_masks
import cv2
from torchvision.transforms.functional import convert_image_dtype


# generate results file
@torch.no_grad()
def generate_results(model, data_loader, device):
    model.eval()

    for i, (image, target) in enumerate(data_loader):

        outputs = model(image)
        img_masks = outputs['masks'] > 0.8
        img_masks = img_masks.squeeze(1)

        img = convert_image_dtype(image, dtype=torch.uint8)
        img = img.squeeze(0)

        # 绘制原图
        write_img = img.numpy().transpose(1, 2, 0)
        write_img = cv2.cvtColor(write_img, cv2.COLOR_BGR2RGB)
        cv2.imwrite("{}.jpg".format(i + 1), write_img)

        # 绘制与之对应的mask图
        tensor = draw_segmentation_masks(img, img_masks, alpha=0.9)
        tensor = tensor.numpy().transpose(1, 2, 0)
        tensor = cv2.cvtColor(tensor, cv2.COLOR_BGR2RGB)
        cv2.imwrite("{}_mask.jpg".format(i+1), tensor)
